#!/usr/bin/python

import pickle
import sys
from copy import deepcopy
from pathlib import Path

import numpy as np

from app.core.config import BILBAO_DATA

from .dealfort import check_input

var = "??"


def read_SiteK(nsg):
    filename = BILBAO_DATA / "BRlist2_A" / f"SiteK_{nsg}.cht"

    # 使用 with 语句自动管理文件资源
    with filename.open("r") as f:
        Nsite = int(f.readline().strip().split()[-1])
        Nktp = int(f.readline().strip().split()[-1])
        # 预分配列表
        name_siteg = [""] * Nsite
        name_irr_siteg = [[] for _ in range(Nsite)]
        name_irr_littg = [[] for _ in range(Nktp)]

        f.readline()  # 跳过空行

        BRdata = [[[] for _ in range(Nktp)] for _ in range(Nsite)]

        for isite in range(Nsite):
            for ikvec in range(Nktp):
                # 读取站点信息
                line = f.readline().strip().split()
                name_siteg[isite] = line[-2]
                nirr_site = int(line[-1])

                # 读取k向量不可约表示数量
                nirr_kvec = int(f.readline().strip().split()[-1])

                # 预分配列表
                name_irr_siteg[isite] = [""] * nirr_site
                name_irr_littg[ikvec] = [""] * nirr_kvec

                # 读取不可约表示名称行
                irr_line = f.readline()
                for i in range(nirr_kvec):
                    name_irr_littg[ikvec][i] = irr_line[7 + 5 * i : 12 + 5 * i].strip()

                # 读取每个不可约表示的数据
                for irr_site in range(nirr_site):
                    data_line = f.readline().strip().split()
                    name_irr_siteg[isite][irr_site] = data_line[0]
                    data = [int(x) for x in data_line[1:]]  # 使用列表推导式
                    BRdata[isite][ikvec].append(data)

                f.readline()  # 跳过空行

    return BRdata, name_siteg, name_irr_siteg, name_irr_littg


def writecode(A, b, lowbound, uppbound, isEBR):
    num_var = len(A[0])
    with open("runZ3.py", "w") as f:
        # f.write("from z3 import Solver, Int, Not, And, sat\n")
        f.write("from z3 import *\n")
        f.write("\n")
        f.write("def decomBR():\n")
        f.write("    s = Solver()\n")
        for i in range(num_var):
            f.write("    a{} = Int('a{}')\n".format(i, i))
        f.write("    vars = [")
        for i in range(num_var):
            f.write("a{}".format(i))
            if i != num_var - 1:
                f.write(",")
        f.write("]\n")

        f.write("\n")
        for i in range(num_var):
            f.write("    s.add(a{} >= {})\n".format(i, lowbound[i]))
            if not isEBR:
                f.write("    s.add(a{} <= {})\n".format(i, uppbound[i]))
        for ieq, eq in enumerate(A):
            tmp_str = ""
            for icoeff, coeff in enumerate(eq):
                tmp_str = tmp_str + "{}*a{}+".format(coeff, icoeff)
            tmp_str = tmp_str + "0 == {}".format(b[ieq])
            f.write("    s.add(" + tmp_str + ")\n")
        # f.write(f"solve({tmp_str})")
        f.write("\n")
        f.write("    alldata = []\n")
        f.write("    #f = open('decompose','a+')\n")
        f.write("    while s.check() == sat:\n")
        f.write("        m = s.model()\n")
        f.write("        if not m:\n")
        # f.write("            print('No solutions any more')\n")
        f.write("            break\n")
        # f.write("        print(m)\n")
        # f.write("    tmp_str = '          '.join([str(m[i]) for i in vars])\n")
        # f.write("    f.write('    ')\n")
        f.write("        tmp_str = '   '.join([str(m[i]) for i in vars])\n")
        f.write("        data = list(map(int,tmp_str.strip().split()))\n")
        f.write("        alldata.append(data)\n")
        f.write("        if len(alldata) >= 8:\n")
        f.write("            break\n")
        f.write("        #for i in range(len(data)):\n")
        f.write("        #    f.write('{:>7d}'.format(data[i]))\n")
        f.write("        s.add(Not(And([v() == m[v] for v in m])))\n")
        # f.write("    else:\n")
        # f.write("        print('No solutions any more')\n")
        f.write("    #f.close()\n")
        f.write("\n")
        f.write("    return alldata\n")
        f.write("\n")
        f.write("\n")
        f.write("if __name__=='__main__':\n")
        f.write("    alldata = decomBR()\n")


def safe_import_runZ3():
    cwd = Path.cwd()

    # 1. 删除 pyc 缓存文件（避免旧版本被加载）
    pycache_path = cwd / "__pycache__"
    if pycache_path.exists():
        for file in pycache_path.glob("runZ3.*.pyc"):
            try:
                file.unlink()
            except Exception as e:
                print(f"⚠️ 删除 {file} 失败: {e}")

    # 2. 删除直接生成的 runZ3.pyc（兼容非 __pycache__ 情况）
    pyc_direct = cwd / "runZ3.pyc"
    if pyc_direct.exists():
        try:
            pyc_direct.unlink()
        except Exception as e:
            print(f"⚠️ 删除 runZ3.pyc 失败: {e}")

    # 3. 插入路径 & 导入
    sys.path.insert(0, str(cwd))
    try:
        sys.modules.pop("runZ3", None)  # 清除缓存模块
        import runZ3  # type: ignore

        return runZ3.decomBR()
    finally:
        if sys.path and sys.path[0] == str(cwd):
            sys.path.pop(0)


def buildAb_oneshot(inputfile, BRdata, soc=0):
    inputfile = Path(inputfile)
    atomwp_path = Path("atomwp.in")
    result_path = inputfile.with_name(f"decompose_{inputfile.name}")

    isEBR = 0

    BR = []
    BRnum = []
    BRset = []
    BRlabel = []
    BRlabelset = []
    with atomwp_path.open("r") as f:
        isEBR = int(f.readline().strip())
        nBR = int(f.readline().strip())
        for i in range(nBR):
            line = f.readline().strip().split()
            BR.append(tuple(map(int, line[0:2])))
            BRlabel.append(line[2])
        for idx, iBR in enumerate(BR):
            if iBR not in BRset:
                BRset.append(iBR)
                BRlabelset.append(BRlabel[idx])
        for i in BRset:
            BRnum.append(BR.count(i))

    maxk = []
    irre_onk = []

    with inputfile.open("r") as f:
        (nsg, nmaxk, nband) = tuple(map(int, f.readline().strip().split()))
        for i in range(nmaxk):
            data = list(map(int, f.readline().strip().split()))
            maxk.append(data[0])
            irre_onk.append(data[1:])

    nirr_eachk = []
    kname = []
    with (BILBAO_DATA / "kvec_list_A.txt").open("r") as f:
        while True:
            (isg, numk) = tuple(map(int, f.readline().strip().split()))
            if isg == nsg:
                for i in range(numk):
                    tmp = f.readline().strip().split()
                    nirr_eachk.append(tuple(map(int, tmp[0:4])))
                    kname.append(tmp[5])
                break
            else:
                for i in range(numk):
                    f.readline()

    for ik in range(len(maxk)):
        if maxk[ik] < 0:
            absk = abs(maxk[ik])
            abskname = kname[absk - 1]
            knameA = abskname + "A"
            kaindex = kname.index(knameA) + 1
            maxk[ik] = kaindex

    nkir = []
    for ik in maxk:
        for jdata in nirr_eachk:
            if ik == jdata[0]:
                if soc == 0:
                    nkir.append(jdata[2])
                else:
                    nkir.append(jdata[3])

    # BRdata, name_siteg, name_irr_siteg, name_irr_littg  = read_SiteK(nsg)

    # construct A matrix
    A = np.zeros((sum(nkir), len(BRset)), dtype=int)

    for iw in range(len(BRset)):
        isite = BRset[iw][0]
        iirre = BRset[iw][1]
        icol = iw
        irow = 0
        for ik in range(len(maxk)):
            for ir in range(nkir[ik]):
                irk = ir if soc == 0 else (ir - nkir[ik])
                A[irow, icol] = A[irow, icol] + BRdata[isite - 1][maxk[ik] - 1][iirre - 1][irk]
                irow = irow + 1

    # construct b array
    b = np.array([0 for i in range(sum(nkir))])
    for ik in range(len(maxk)):
        for rep in irre_onk[ik]:
            if soc != 0:
                rep = rep - nirr_eachk[maxk[ik] - 1][2]
            b[sum(nkir[:ik]) + rep - 1] += 1

    lowbound = np.array([0 for i in range(len(A[0]))])
    uppbound = np.array(BRnum)

    writecode(A, b, lowbound, uppbound, isEBR)

    # 将当前工作目录（即 runZ3.py 所在的目录）添加到 sys.path
    alldata = safe_import_runZ3()
    num_decom = len(alldata)
    alldata = np.array(alldata).T

    with result_path.open("w") as f:
        f.write("{:>7d}\n".format(num_decom))
        for idx, ibr in enumerate(BRset):
            f.write(
                "{:>4d}".format(idx + 1)
                + "  {:>2d}@{:<2d}".format(ibr[1], ibr[0])
                + "{:>10}".format(BRlabelset[idx])
                + "    ({:3d}) :".format(uppbound[idx])
            )
            if len(alldata) != 0:
                for jdx, jbr in enumerate(alldata[idx]):
                    f.write("{:>3d};".format(alldata[idx][jdx]))
            f.write("\n")

    return num_decom


def run_buildAb(inputfile="fort.67", soc=0):
    inputfile = Path(inputfile)
    nsg, num_vars, k_withvar = check_input(inputfile)
    BRdata, *_ = read_SiteK(nsg)
    nirr_eachk = []
    with open(BILBAO_DATA / "kvec_list_A.txt", "r") as f:
        while True:
            (isg, numk) = tuple(map(int, f.readline().strip().split()))
            if isg == nsg:
                for i in range(numk):
                    nirr_eachk.append(tuple(map(int, f.readline().strip().split()[0:4])))
                break
            else:
                for i in range(numk):
                    f.readline()

    irr_trial = []
    for ik in range(len(k_withvar)):
        if soc == 0:
            irr_trial.append([i + 1 for i in range(nirr_eachk[k_withvar[ik] - 1][2])])
        else:
            irr_trial.append(
                [nirr_eachk[k_withvar[ik] - 1][2] + i + 1 for i in range(nirr_eachk[k_withvar[ik] - 1][3])]
            )

    summary_file = Path("solutions")
    if num_vars == 0:
        num_decom = buildAb_oneshot(inputfile, BRdata, soc)
    elif num_vars == 1:
        with summary_file.open("w") as summary:
            summary.write("# trial_irr  num_solutions\n")
            for irr1 in irr_trial[0]:
                new_inputfile = write_input(inputfile, [irr1])
                num_decom = buildAb_oneshot(new_inputfile, BRdata, soc)
                if num_decom > 0:
                    summary.write("{:>5d}:{:>5d}\n".format(irr1, num_decom))
                else:
                    summary.write("{:>5d}:\n".format(irr1))
    elif num_vars == 2:
        with summary_file.open("w") as summary:
            summary.write("# trial_irr1  trial_irr2  num_solutions\n")
            for irr1 in irr_trial[0]:
                for irr2 in irr_trial[1]:
                    new_inputfile = write_input(inputfile, [irr1, irr2])
                    num_decom = buildAb_oneshot(new_inputfile, BRdata, soc)
                    if num_decom > 0:
                        summary.write("{:>5d}{:>5d}:{:>5d}\n".format(irr1, irr2, num_decom))
                    else:
                        summary.write("{:>5d}{:>5d}:\n".format(irr1, irr2))

    elif num_vars == 3:
        with summary_file.open("w") as summary:
            summary.write("# trial_irr1  trial_irr2  trial_irr3  num_solutions\n")
            for irr1 in irr_trial[0]:
                for irr2 in irr_trial[1]:
                    for irr3 in irr_trial[2]:
                        new_inputfile = write_input(inputfile, [irr1, irr2, irr3])
                        num_decom = buildAb_oneshot(new_inputfile, BRdata, soc)
                        if num_decom > 0:
                            summary.write("{:>5d}{:>5d}{:>5d}:{:>5d}\n".format(irr1, irr2, irr3, num_decom))
                        else:
                            summary.write("{:>5d}{:>5d}{:>5d}:\n".format(irr1, irr2, irr3))
    else:
        raise ValueError(f"{num_vars}(>3) variables in input file")

    return


def write_input(inputfile, irr_trial: list[int]) -> Path:
    inputfile = Path(inputfile)
    # 构造新文件名
    newfile = inputfile.with_name(f"{inputfile.name}_{'_'.join(map(str, irr_trial))}")

    with inputfile.open("r") as f, newfile.open("w") as w:
        tmp = f.readline()
        numk = int(tmp.split()[1])
        w.write(tmp)

        for ik in range(numk):
            tmp = f.readline().strip().split()
            for d in tmp:
                if d == var:
                    w.write("{:>3d}".format(irr_trial.pop(0)))
                else:
                    w.write("{:>3s}".format(d))
            w.write("\n")

    return newfile


def read_rsi(filename: Path) -> dict:
    with filename.open("r") as f:
        fdata = f.readlines()

    data = dict()
    i = 0
    while i < len(fdata) - 1:
        w = fdata[i].strip()
        data[w] = dict()
        i += 1

        data[w]["row"] = []
        i += 1

        while fdata[i].strip().split()[0] != "Column":
            data[w]["row"].append(fdata[i].strip())
            i += 1
        lenrow = len(data[w]["row"])

        data[w]["col"] = []
        i += 1

        while fdata[i].strip().split()[0] != "C":
            data[w]["col"].append(fdata[i].strip())
            i += 1
        lencol = len(data[w]["col"])

        data[w]["rep"] = np.zeros((lenrow, lencol), dtype=int)
        i += 1

        for count in range(lenrow):
            data[w]["rep"][count] = np.array(list(map(int, fdata[i].strip().split())))
            i += 1

        data[w]["ind"] = np.zeros((lenrow, lencol), dtype=int)
        i += 1

        for count in range(lenrow):
            data[w]["ind"][count] = np.array(list(map(int, fdata[i].strip().split())))
            i += 1

        data[w]["L"] = np.zeros((lenrow, lenrow), dtype=int)
        i += 1

        for count in range(lenrow):
            data[w]["L"][count] = np.array(list(map(int, fdata[i].strip().split())))
            i += 1

        i += 1
        i += 1
    return data


def run_rsi(decompose_ebr="decompose.601", brdatafile="fort.67", soc=0, trs=1):
    # 转换为Path对象
    decompose_ebr = Path(decompose_ebr)
    brdatafile = Path(brdatafile)
    output_path = Path("rsidelta")

    # 使用Path对象和with语句处理数据文件
    pgirr_path = BILBAO_DATA / "pgirr.pkl"
    wp_list_path = BILBAO_DATA / "wp_list.pkl"

    with pgirr_path.open("rb") as f:
        replist = pickle.load(f)
    with wp_list_path.open("rb") as f:
        wplist = pickle.load(f)

    nsg, num_vars, k_withvar = check_input(brdatafile)
    if soc == 0 and trs == 1:
        rsi_data_path = BILBAO_DATA / "src_rsi" / "nosoc_trs" / f"{nsg}.txt"
        data = read_rsi(rsi_data_path)
    else:
        raise NotImplementedError("RSI calculation is only supported for soc=0 and trs=1.")

    # 使用with语句和Path对象读取分解文件
    if not decompose_ebr.exists():
        raise FileNotFoundError(f"分解文件 {decompose_ebr} 未找到")

    fdata = decompose_ebr.read_text().splitlines()
    num_decom = int(fdata[0].strip())
    if num_decom == 0:
        raise ValueError(f"No decomposition in the file {str(decompose_ebr)}")

    decomT = []
    ebrname = []
    ebrname_str = []
    ebrname_ind = []
    for line in fdata[1:]:
        tmp0 = line.strip().split(":")[0]
        tmp = line.strip().split(":")[1]
        decomT.append(list(map(int, tmp.strip().split(";")[:-1])))
        ebrname.append(tmp0.strip().split()[1])
        ebrname_str.append(tmp0.strip().split()[2])
        ebrname_ind.append(tmp0.strip().split()[0])
    decomT = np.array(decomT)

    ebrname_GM = []
    for ebr in ebrname_str:
        nameirr = ebr.split("@")[0]
        w = ebr.split("@")[1]
        pgname = wplist[nsg][w]

        for pg in replist[1:]:
            if pg["name"] == pgname:
                break
        else:
            raise ValueError(f"Point group '{pgname}' not found in representation list for ebr '{ebr}'.")

        for ind, irname in enumerate(pg["rep3"]):
            if irname == nameirr:
                break
        else:
            raise ValueError(f"Irreducible representation '{nameirr}' not found in point group '{pgname}'.")

        ebrname_GM.append(pg["rep1"][ind] + "@" + w)

    for w in data.keys():
        data[w]["row_A"] = []

        pgname = wplist[nsg][w]
        for pg in replist[1:]:
            if pg["name"] == pgname:
                break
        else:
            raise ValueError(f"Point group '{pgname}' not found for Wyckoff position '{w}'.")
        pgtmp = deepcopy(pg)
        pgtmp["rep1"].reverse()
        pgtmp["rep3"].reverse()
        GMlist = pgtmp["rep1"]
        Alist = pgtmp["rep3"]
        for rep in data[w]["row"]:
            reptmp = rep
            for repGM, repA in zip(GMlist, Alist):
                reptmp = reptmp.replace(repGM, repA)
            data[w]["row_A"].append(reptmp)

    existwyck = []
    for ebr in ebrname_GM:
        w = ebr.split("@")[1]
        if w not in existwyck:
            existwyck.append(w)

    # 使用with语句和Path对象写入文件
    with output_path.open("w") as wfile:
        wfile.write("{:>4d}\n".format(len(existwyck)))

        for w in existwyck:
            wfile.write("{:>5s}".format(w))
            if w not in data.keys():
                raise KeyError(f"wyck在essebr {w:>5s} 中不在rsi公式中")
            lenrow = len(data[w]["row"])
            lencol = len(data[w]["col"])

            constraint = []
            ebrind = []
            for ebr in data[w]["row"]:
                if "&" in ebr:
                    e1 = ebr.split("&")[0]
                    e2 = ebr.split("&")[1]
                    ebrind.append(ebrname_GM.index(e1))
                    constraint.append([ebrname_GM.index(e1), ebrname_GM.index(e2)])
                else:
                    ebrind.append(ebrname_GM.index(ebr))

            decomW = np.zeros((lenrow, num_decom), dtype=int)
            for isol in range(num_decom):
                for constr in constraint:
                    c1 = constr[0]
                    c2 = constr[1]
                    if decomT[c1, isol] != decomT[c2, isol]:
                        decomW[:, isol] = 0
                        break
                else:
                    for ii, ind in enumerate(ebrind):
                        decomW[ii, isol] = decomT[ind, isol]

            delta = np.dot(data[w]["L"], decomW)

            module = []
            for irank in range(lenrow):
                if irank < min(lenrow, lencol):
                    module.append(data[w]["ind"][irank, irank])
                else:
                    module.append(0)

            wfile.write("{:5d}{:5d}{:5d}\n".format(lenrow, num_decom, len(module) - module.count(1)))
            for imod in module:
                wfile.write("{:>4d}".format(imod))
            wfile.write("\n")

            # rsi formula
            for ind, imod in enumerate(module):
                if imod != 1:
                    for icol in range(lenrow):
                        if data[w]["L"][ind, icol] != 0:
                            ebrname_A = data[w]["row_A"][icol]
                            if data[w]["L"][ind, icol] == 1:
                                wfile.write("+")
                            elif data[w]["L"][ind, icol] == -1:
                                wfile.write("-")
                            elif data[w]["L"][ind, icol] > 0:
                                wfile.write("+" + str(data[w]["L"][ind, icol]))
                            else:
                                wfile.write(str(data[w]["L"][ind, icol]))
                            wfile.write(r"n(" + ebrname_A + ")")
                    wfile.write("\n")

            for irow in range(lenrow):
                for icol in range(num_decom):
                    wfile.write("{:>3d}".format(delta[irow, icol]))
                wfile.write("\n")

    return


def run_essebr(
    decompose_ebr="decompose.601",
    decompose_abr="decompose.603",
    brdatafile="fort.67",
    soc=0,
):
    decompose_ebr = Path(decompose_ebr)
    decompose_abr = Path(decompose_abr)
    brdatafile = Path(brdatafile)

    if not decompose_ebr.exists():
        raise FileNotFoundError(f"File {decompose_ebr} not found.")

    fdata = decompose_ebr.read_text().splitlines()
    num_decom = int(fdata[0].strip())
    if num_decom == 0:
        raise ValueError(f"No decomposition in the file {decompose_ebr}")

    decom = []
    ebrname = []
    ebrname_str = []
    ebrname_ind = []
    for line in fdata[1:]:
        tmp0, tmp = line.strip().split(":", 1)
        decom.append(list(map(int, tmp.strip().split(";")[:-1])))

        parts = tmp0.strip().split()
        if len(parts) >= 3:
            ebrname_ind.append(parts[0])
            ebrname.append(parts[1])
            ebrname_str.append(parts[2])

    if not decompose_abr.exists():
        raise FileNotFoundError(f"File {decompose_abr} not found.")

    fdata = decompose_abr.read_text().splitlines()
    num_decom = int(fdata[0].strip())
    if num_decom > 0:
        raise ValueError(f"Decompositions in the file {decompose_abr} already exist.")

    essebr_list = ebrname
    essebr_str_list = ebrname_str
    essebr_ind_list = ebrname_ind

    nsg, num_vars, k_withvar = check_input(brdatafile)
    BRdata, *_ = read_SiteK(nsg)

    essebr_path = Path("essebr")

    with essebr_path.open("w") as essw:
        for ind, essebr in enumerate(essebr_list):
            new_inputfile = write_inputrmebr(brdatafile, essebr, BRdata)
            if new_inputfile is None:
                continue
            num_decom = buildAb_oneshot(new_inputfile, BRdata, soc)
            essw.write(
                "{:>4s} {:>10s} {:>10s} {:>10d}\n".format(essebr_ind_list[ind], essebr, essebr_str_list[ind], num_decom)
            )

    return


def write_inputrmebr(inputfile, essebr, BRdata):
    inputfile = Path(inputfile)
    nsg, num_vars, k_withvar = check_input(inputfile)

    nirr_eachk = []
    kname = []
    kvec_list_file = BILBAO_DATA / "kvec_list_A.txt"
    with kvec_list_file.open("r") as f:
        while True:
            (isg, numk) = tuple(map(int, f.readline().strip().split()))
            if isg == nsg:
                for i in range(numk):
                    tmp = f.readline().strip().split()
                    nirr_eachk.append(tuple(map(int, tmp[0:4])))
                    kname.append(tmp[5])
                break
            else:
                for i in range(numk):
                    f.readline()

    ess_site = int(essebr.strip().split("@")[1])
    ess_irr = int(essebr.strip().split("@")[0])

    removedict = dict()
    for ik in range(len(BRdata[0])):
        removedict[ik + 1] = []
        for ind, ir in enumerate(BRdata[ess_site - 1][ik][ess_irr - 1]):
            for i in range(ir):
                removedict[ik + 1].append(ind + 1)

    maxklist = []

    newfile = inputfile.with_name(f"{inputfile.name}_{essebr.replace('@', 'a')}")

    with inputfile.open("r") as f, newfile.open("w") as w:
        tmp = f.readline()
        numk = int(tmp.split()[1])
        w.write(tmp)
        for ik in range(numk):
            tmp = list(map(int, f.readline().strip().split()))
            tmpir = tmp[1:]
            indk = tmp[0]
            if indk < 0:
                abskname = kname[abs(indk) - 1]
                knameA = abskname + "A"
                indk = kname.index(knameA) + 1
            maxklist.append(indk)
            for ir in removedict[indk]:
                try:
                    tmpir.remove(ir)
                except ValueError:
                    return None
            w.write("{:>3d}".format(indk))
            for r in tmpir:
                w.write("{:>3d}".format(r))
            w.write("\n")
        w.write("\n")
        w.write(essebr + "\n")
        for key in removedict:
            if key in maxklist:
                w.write("{:>3d}".format(key))
                for r in removedict[key]:
                    w.write("{:>3d}".format(r))
                w.write("\n")
    return newfile
